from flask import Flask, request, jsonify, send_from_directory
from PIL import Image
import io
import base64
import os

app = Flask(__name__)

# 文本转二进制
def text_to_binary(text):
    return ''.join(format(ord(char), '08b') for char in text)

# 二进制转文本
def binary_to_text(binary):
    text = ''
    for i in range(0, len(binary), 8):
        byte = binary[i:i+8]
        if byte:
            text += chr(int(byte, 2))
    return text

# 计算图片最大容量（以字节为单位）
def calculate_max_capacity(image):
    width, height = image.size
    total_pixels = width * height
    # 每个像素3个颜色通道，每个通道可以存储1bit，减去结束标记(1字节)
    return (total_pixels * 3) // 8 - 1

# 嵌入信息到图片
def embed_text(image, text):
    # 转换为RGB模式
    image = image.convert('RGB')
    width, height = image.size

    # 计算容量
    max_bytes = calculate_max_capacity(image)

    # 编码为UTF-8字节并添加结束标记
    utf8_bytes = text.encode('utf-8')
    if len(utf8_bytes) > max_bytes:
        raise ValueError(f"信息过长，图片最大可存储 {max_bytes} 个字节（约 {max_bytes // 3} 个中文）")

    # 添加结束标记
    utf8_bytes += b'\x00'

    # 转换为二进制
    binary_text = ''.join(format(byte, '08b') for byte in utf8_bytes)
    binary_length = len(binary_text)

    # 嵌入到像素中
    data_index = 0
    for y in range(height):
        for x in range(width):
            r, g, b = image.getpixel((x, y))

            # 处理红色通道
            if data_index < binary_length:
                r = (r & 0xFE) | int(binary_text[data_index])
                data_index += 1
            # 处理绿色通道
            if data_index < binary_length:
                g = (g & 0xFE) | int(binary_text[data_index])
                data_index += 1
            # 处理蓝色通道
            if data_index < binary_length:
                b = (b & 0xFE) | int(binary_text[data_index])
                data_index += 1

            image.putpixel((x, y), (r, g, b))

            # 完成嵌入
            if data_index >= binary_length:
                return image
    return image

# 从图片提取信息
def extract_text(image):
    # 转换为RGB模式
    image = image.convert('RGB')
    width, height = image.size

    binary_text = ''
    for y in range(height):
        for x in range(width):
            r, g, b = image.getpixel((x, y))

            # 提取红色通道最低位
            binary_text += str(r & 0x01)
            # 提取绿色通道最低位
            binary_text += str(g & 0x01)
            # 提取蓝色通道最低位
            binary_text += str(b & 0x01)

    # 转换为字节列表
    bytes_list = []
    for i in range(0, len(binary_text), 8):
        byte_str = binary_text[i:i+8]
        if len(byte_str) < 8:
            break  # 不足8位，结束
        byte = int(byte_str, 2)
        if byte == 0x00:
            break  # 遇到结束标记，结束
        bytes_list.append(byte)

    try:
        # 解码为UTF-8文本
        text = bytes(bytes_list).decode('utf-8')
        return text
    except UnicodeDecodeError as e:
        raise ValueError(f"提取失败：{str(e)}")

# 主页面路由
@app.route('/')
def index():
    return send_from_directory('.', 'index.html')

# 嵌入路由
@app.route('/embed', methods=['POST'])
def embed():
    try:
        # 检查文件
        if 'image' not in request.files:
            return jsonify({'error': '未上传图片'}), 400

        image_file = request.files['image']
        if image_file.filename == '':
            return jsonify({'error': '未选择图片'}), 400

        # 检查文件类型
        if image_file.content_type not in ['image/png', 'image/jpeg']:
            return jsonify({'error': '仅支持PNG/JPG格式图片'}), 400

        # 检查文件大小
        if image_file.content_length > 5 * 1024 * 1024:
            return jsonify({'error': '图片大小不能超过5MB'}), 400

        # 读取图片
        image = Image.open(image_file)

        # 检查文本
        text = request.form.get('text', '').strip()
        if not text:
            return jsonify({'error': '请输入要隐藏的文本'}), 400

        # 嵌入文本
        stego_image = embed_text(image, text)

        # 保存为PNG格式
        img_buffer = io.BytesIO()
        stego_image.save(img_buffer, format='PNG')
        img_buffer.seek(0)

        return app.response_class(img_buffer.read(), mimetype='image/png')

    except ValueError as e:
        return jsonify({'error': str(e)}), 400
    except Exception as e:
        return jsonify({'error': f'操作失败：{str(e)}'}), 500

# 提取路由
@app.route('/extract', methods=['POST'])
def extract():
    try:
        # 检查文件
        if 'image' not in request.files:
            return jsonify({'error': '未上传图片'}), 400

        image_file = request.files['image']
        if image_file.filename == '':
            return jsonify({'error': '未选择图片'}), 400

        # 检查文件类型
        if image_file.content_type not in ['image/png', 'image/jpeg']:
            return jsonify({'error': '仅支持PNG/JPG格式图片'}), 400

        # 读取图片
        image = Image.open(image_file)

        # 提取文本
        text = extract_text(image)

        if not text:
            return jsonify({'error': '未检测到隐藏信息'}), 400

        return jsonify({'text': text})

    except Exception as e:
        return jsonify({'error': f'操作失败：{str(e)}'}), 500

if __name__ == '__main__':
    app.run(debug=True)
